﻿//
// Created by QU on 24-10-6.
//

#ifndef WORKER_POOL_H
#define WORKER_POOL_H

#include <boost/asio.hpp>
#include <vector>
#include <thread>
#include <memory>
#include <atomic>
#include <mutex>

#include "server/server/Worker.h"
#include "server/public/Singleton.hpp"

using tcp = asio::ip::tcp;
namespace asio = boost::asio;

class WorkerPool : public Singleton<WorkerPool>
{
    friend class Singleton<WorkerPool>;

public:
    WorkerPool(const WorkerPool&) = delete;
    WorkerPool& operator=(const WorkerPool&) = delete;
    WorkerPool(WorkerPool&&) = delete;

    bool startAcceptor(uint16_t port);
    bool stopAcceptor();

    bool startWorkers(size_t worker_thread_count);
    bool requestStopWorkers();
    bool forceStopWorkers();

    bool startAll(uint16_t port, size_t worker_count = 2);
    void requestStopAll();
    void forceStopAll();

    [[nodiscard]] asio::io_context& getAcceptorIOContext();

private:
    WorkerPool();
    ~WorkerPool();

    void handleAccept(const boost::system::error_code& error, tcp::socket&& socket);
    void doAccept(std::stop_token& stoken);
    static void printRemoteSockInformation(const ip::tcp::socket& sock);

    size_t getActiveWorkerCount() const;
    std::vector<size_t> getWorkerSessionCounts() const;
    Worker& getLeastLoadedWorker();
    size_t getTotalSessionCount() const;

    asio::io_context m_acceptor_io_context;
    asio::executor_work_guard<asio::io_context::executor_type> m_acceptor_io_work_guard;
    std::unique_ptr<tcp::acceptor> m_acceptor;
    std::jthread m_acceptor_thread;

    std::shared_ptr<std::stop_source> m_shared_worker_stop_source;
    std::vector<std::unique_ptr<asio::io_context>> m_worker_io_context_s;
    std::vector<std::unique_ptr<Worker>> m_worker_s;

    mutable std::mutex m_mutex;
};

#endif // WORKER_POOL_H
